#include <windows.h>
#include <map>
#include <iostream>

#include <list>


//#define TSFUNCS_EXCLUDE_MINHOOK /* This program doesn't use any hooks */

#include "MinHook/MinHook.h"

#include "TSfuncs/TSFuncs.cpp"
#include "TSfuncs/TSHooks.cpp"
#include "TSfuncs/TSVector.cpp"

#include "BLstructs/include.hpp"
#include "main.hpp"




std::map<ADDR*, alloc_struct*> mem_map;
bool journaling_active = false;
unsigned int LIMIT_LINES = 12500;
std::list<std::string> object_journal;


int ObjectTrace_debug_memory = 0;


BlFunctionDef(ADDR*, __fastcall, op__new, unsigned int, void*);
BlFunctionHookDef(op__new);

ADDR* __fastcall op__newHook(unsigned int bytes, void* edx)
{
	ADDR* loc = op__newOriginal(bytes, edx);
	if(loc != NULL)
	{
		if(mem_map[loc] == NULL)
			mem_map[loc] = new alloc_struct { loc, bytes, 0, 0, 0, 0, 0, 0 };

		else
			mem_map[loc]->length = bytes;
	}

	if(ObjectTrace_debug_memory)
		printf("Alloc: %p, %d bytes\n", loc, bytes);
		
	if(journaling_active)
		journal_event(string_format("0x%p\tALLOC\tlength=%d", loc, bytes));

	return loc;
}




BlFunctionDef(void, __fastcall, op__delete, ADDR*, void*);
BlFunctionHookDef(op__delete);

void __fastcall op__deleteHook(ADDR* loc, void* edx)
{
	if(ObjectTrace_debug_memory)
		printf("--Dealloc: %p\n", loc);
	
	//NULL check not necessary
	if(mem_map[loc] != NULL)
	{
		//printf("UNLOAD LOC %p, %p\n", loc, mem_map[loc]);
		delete mem_map[loc];
		mem_map.erase(loc);
		//printf("   Next %p\n", mem_map[loc]);
	}

	if(journaling_active)
		journal_event(string_format("0x%p\tDEALLOC", loc));

	op__deleteOriginal(loc, edx);
}




int ObjectTrace_debug_objects = 0;


BlFunctionDef(bool, __fastcall, SimObject__onAdd, ADDR*, void*);
BlFunctionHookDef(SimObject__onAdd);

bool __fastcall SimObject__onAddHook(ADDR* objADDR, void *edx)
{
	bool result = SimObject__onAddOriginal(objADDR, edx);

	if(mem_map[objADDR] != NULL)
	{
		mem_map[objADDR]->isSimObject = true;
	}

	if(ObjectTrace_debug_objects)
		printf(":::CONSTRUCTED %s\n", RTTI_className(objADDR));

	if(journaling_active)
		journal_event(string_format("0x%p\tSimObject::onAdd\t%s", objADDR, RTTI_className(objADDR)));

	return result;
}


BlFunctionDef(bool, __fastcall, SimObject__registerObject, ADDR*, void*);
BlFunctionHookDef(SimObject__registerObject);

bool __fastcall SimObject__registerObjectHook(ADDR* objADDR, void *edx)
{
	bool result = SimObject__registerObjectOriginal(objADDR, edx);

	bool seen_before = false;
	if(mem_map[objADDR] != NULL)
	{
		seen_before = true;

		if(result)
			mem_map[objADDR]->SimObjectRegistered = true;
		else
			mem_map[objADDR]->SimObjectRegisterFailure = true;

		
		mem_map[objADDR]->simID = ((SimObject*)objADDR)->id;
	}

	if(journaling_active)
		journal_event(string_format("0x%p\tSimObject::registerObject\tresult=%d\tseen=%d", objADDR, result, seen_before));

	return result;
}


BlFunctionDef(void, __fastcall, SimObject__onRemove, ADDR*, void*);
BlFunctionHookDef(SimObject__onRemove);

void __fastcall SimObject__onRemoveHook(ADDR* objADDR, void *edx)
{
	if(ObjectTrace_debug_objects)
		printf(":::DESTRUCTED %s\n", RTTI_className(objADDR));
		
	if(mem_map[objADDR] != NULL)
	{
		mem_map[objADDR]->SimObjectRemoved = true;
	}

	if(journaling_active)
		journal_event(string_format("0x%p\tSimObject::onRemove", objADDR));

	SimObject__onRemoveOriginal(objADDR, edx);
}





void TS_trace(ADDR obj, int argc, const char *argv[])
{
	bool lastEnabled = journaling_active;
	journaling_active = atoi(argv[1]) || stricmp(argv[1], "true") == 0;

	if(lastEnabled != journaling_active)
	{
		if(journaling_active)
			BlPrintf("ObjectTrace is enabled.");
		else
			BlPrintf("ObjectTrace is disabled.");
	}

	else
	{
		if(journaling_active)
			BlPrintf("ObjectTrace is already enabled.");
		else
			BlPrintf("ObjectTrace is already disabled.");
	}
}

int TS_getLineCount(ADDR obj, int argc, const char *argv[])
{
	return object_journal.size();
}

int TS_getLineLimit(ADDR obj, int argc, const char *argv[])
{
	return LIMIT_LINES;
}

void TS_setLineLimit(ADDR obj, int argc, const char *argv[])
{
	int new_limit = atoi(argv[1]);
	if(new_limit <= 0)
	{
		BlPrintf("objectJournalSetLineLimit() - can't have limit less than 1");
		return;
	}
	
	LIMIT_LINES = new_limit;
}

int TS_popLines(ADDR obj, int argc, const char *argv[])
{
	int lines_to_remove = atoi(argv[1]);
	if(lines_to_remove <= 0)
	{
		BlPrintf("objectJournalPopLines() - can't remove less than one line");
		return 0;
	}

	int lines_removed = 0;
	while(object_journal.size() > 0 && lines_removed <= lines_to_remove)
	{
		object_journal.pop_front();
		lines_removed++;
	}
	
	return lines_removed;
}

int TS_clear(ADDR obj, int argc, const char *argv[])
{
	int lines_removed = object_journal.size();
	while(object_journal.size() > 0)
		object_journal.pop_front();

	return lines_removed;
}

int TS_print_journal(ADDR obj, int argc, const char *argv[])
{
	int lines_printed = 0;
	for(auto const& i : object_journal)
	{
		BlPrintf(std::string(i).c_str());
		lines_printed++;
	}

	return lines_printed;
}

void journal_event(std::string event_text)
{
	if(!journaling_active)
		return;

	if(object_journal.size() > LIMIT_LINES)
		object_journal.pop_front();
		
	object_journal.push_back(event_text +"\n");
}


bool init()
{
	BlInit;

	if (!tsf_InitInternal())
		return false;


	tsf_AddVar("$ObjectTrace_debug_objects", &ObjectTrace_debug_objects);
	tsf_AddVar("$ObjectTrace_debug_memory", &ObjectTrace_debug_memory);

	//0x578160	op__new
	BlScanFunctionHex(op__new, "55 8b ec 83 e4 f8 80 3d ? ? ? ? 00");
	BlCreateHook(op__new);
	BlTestEnableHook(op__new);

	//0x576e40	op__delete
	BlScanFunctionHex(op__delete, "51 53 55 56 8b 35 ? ? ? ?");
	BlCreateHook(op__delete);
	BlTestEnableHook(op__delete);
	
		
	//0x44a440	SimObject__onAdd
	BlScanFunctionHex(SimObject__onAdd, "56 8b f1 83 4e 18 08");
	BlCreateHook(SimObject__onAdd);
	BlTestEnableHook(SimObject__onAdd);

	//0x44d290	SimObject__registerObject
	BlScanFunctionHex(SimObject__registerObject, "51 53 56 8b f1 83 66 18 fc");
	BlCreateHook(SimObject__registerObject);
	BlTestEnableHook(SimObject__registerObject);

	//0x44a4e0	SimObject__onRemove
	BlScanFunctionHex(SimObject__onRemove, "83 61 18 f7 8b 51 2c");
	BlCreateHook(SimObject__onRemove);
	BlTestEnableHook(SimObject__onRemove);

	tsf_AddConsoleFunc(NULL, NULL, "objectJournal", TS_trace, "(bool enabled)", 2, 2);
	tsf_AddConsoleFunc(NULL, NULL, "objectJournalGetLineCount", TS_getLineCount, "()", 1, 1);
	tsf_AddConsoleFunc(NULL, NULL, "objectJournalGetLineLimit", TS_getLineLimit, "()", 1, 1);
	tsf_AddConsoleFunc(NULL, NULL, "objectJournalSetLineLimit", TS_setLineLimit, "(int size)", 2, 2);
	tsf_AddConsoleFunc(NULL, NULL, "objectJournalPopLines", TS_popLines, "(int amount)", 2, 2);
	tsf_AddConsoleFunc(NULL, NULL, "objectJournalClear", TS_clear, "()", 1, 1);
	tsf_AddConsoleFunc(NULL, NULL, "objectJournalPrintJournal", TS_print_journal, "()", 1, 1);

	BlPrintf("    object_journal pointer: 0x%p", &object_journal);
	
	BlPrintf("%s (v%s-%s): init'd", PROJECT_NAME, PROJECT_VERSION, TSFUNCS_DEBUG ? "debug":"release");

	return true;
}

bool deinit(bool terminating)
{
	printf("%s: deinit'd\n", PROJECT_NAME);
	//if(!terminating)
	//	BlTestDisableHook(CodeBlock__exec);

	if(!journaling_active)
	{
		printf("    - ObjectTrace was disabled\n");
		return true;
	}

	return true;
}

bool __stdcall DllMain(HINSTANCE hinstance, int reason, void *reserved)
{
	switch (reason)
	{
		case DLL_PROCESS_ATTACH:
			return init();
		case DLL_PROCESS_DETACH:
			return deinit(reserved != NULL);
		default:
			return true;
	}
}

extern "C" void __declspec(dllexport) __cdecl PROJECT_EXPORT(){}
